import random

import gym
import numpy as np
from gym.spaces import Discrete, Box


class OOXXGame(gym.Env):
    def __init__(self):
        self.action_space = Discrete(9)
        # 2x2 Box
        self.observation_space = Box(0, 1, (3, 3), dtype=np.float32)
        self._max_episode_steps = 6

    def reset(self, *args, **kwargs):
        self.obs = np.zeros((3, 3))
        self.black_step = True
        return self.obs

    def is_step_black(self):
        black_num = np.where(self.obs == 1)[0]
        white_num = np.where(self.obs == 0.5)[0]
        if len(black_num) > len(white_num):
            return False
        else:
            return True

    def step(self, action):
        if self.black_step:
            self.black_step = False
            return self.step_black(action)
        else:
            self.black_step = True
            return self.step_white(action)

    def step_black(self, action):
        if not self.is_step_black():
            return self.obs, 0, False, True
        reward = 0
        # 将action转为坐标
        x = action // 3
        y = action % 3
        done = True
        # 如果该位置已经有棋子了，reward = -1
        if self.obs[x][y] != 0:
            reward = -1
            done = False
        else:
            # 落子
            self.obs[x][y] = 1
            # 如果黑子排成一行
            if self.obs[0][0] == 1 and self.obs[0][1] == 1 and self.obs[0][2] == 1:
                reward = 1
            elif self.obs[1][0] == 1 and self.obs[1][1] == 1 and self.obs[1][2] == 1:
                reward = 1
            elif self.obs[2][0] == 1 and self.obs[2][1] == 1 and self.obs[2][2] == 1:
                reward = 1
            elif self.obs[0][0] == 1 and self.obs[1][0] == 1 and self.obs[2][0] == 1:
                reward = 1
            elif self.obs[0][1] == 1 and self.obs[1][1] == 1 and self.obs[2][1] == 1:
                reward = 1
            elif self.obs[0][2] == 1 and self.obs[1][2] == 1 and self.obs[2][2] == 1:
                reward = 1
            elif self.obs[0][0] == 1 and self.obs[1][1] == 1 and self.obs[2][2] == 1:
                reward = 1
            elif self.obs[0][2] == 1 and self.obs[1][1] == 1 and self.obs[2][0] == 1:
                reward = 1
            # 如果棋盘已满，reward = -1
            elif np.sum(self.obs) == 5 * 1 + 4 * 0.5:
                reward = -1
            else:
                done = False
        return self.obs, reward, done, True

    def step_white(self, action):
        if self.is_step_black():
            return self.obs, 0, False, True
        reward = 0
        # 将action转为坐标
        x = action // 3
        y = action % 3
        done = True
        if np.sum(self.obs) == 5 * 1 + 4 * 0.5:
            reward = -1
        # 如果该位置已经有棋子了，reward = -1
        elif self.obs[x][y] != 0:
            reward = -1
            done = False
        else:
            # 落子
            self.obs[x][y] = 0.5
            # 如果白子排成一行
            if self.obs[0][0] == 0.5 and self.obs[0][1] == 0.5 and self.obs[0][2] == 0.5:
                reward = 1
            elif self.obs[1][0] == 0.5 and self.obs[1][1] == 0.5 and self.obs[1][2] == 0.5:
                reward = 1
            elif self.obs[2][0] == 0.5 and self.obs[2][1] == 0.5 and self.obs[2][2] == 0.5:
                reward = 1
            elif self.obs[0][0] == 0.5 and self.obs[1][0] == 0.5 and self.obs[2][0] == 0.5:
                reward = 1
            elif self.obs[0][1] == 0.5 and self.obs[1][1] == 0.5 and self.obs[2][1] == 0.5:
                reward = 1
            elif self.obs[0][2] == 0.5 and self.obs[1][2] == 0.5 and self.obs[2][2] == 0.5:
                reward = 1
            elif self.obs[0][0] == 0.5 and self.obs[1][1] == 0.5 and self.obs[2][2] == 0.5:
                reward = 1
            elif self.obs[0][2] == 0.5 and self.obs[1][1] == 0.5 and self.obs[2][0] == 0.5:
                reward = 1
            else:
                done = False
        return self.obs, reward, done, True

    def render(self, *args, **kwargs):
        # 1用0表示，0.5用x表示,未落子的地方用-表示
        for x in range(3):
            for y in range(3):
                if self.obs[x][y] == 0.5:
                    print('X', end='')
                elif self.obs[x][y] == 1:
                    print('O', end='')
                else:
                    print('-', end='')
            print()
        print()


if __name__ == '__main__':
    a = np.array([[1, 1, 0], [1, 0, 1], [1, 0, 0]])
    print(len(np.where(a == 1)[0]))
